#include <regex>
#include "main/utils.h"
using namespace std;

struct PrefixNode {
  bool terminal = 0;
  PrefixNode* childern[26];
};

void insertToPrefixTree(PrefixNode* root, const string& prefix) {
  PrefixNode* node = root;
  for (const char& ele : prefix) {
    if (node->childern[ele - 'a'] == nullptr) {
      node->childern[ele - 'a'] = new PrefixNode();
    }
    node = node->childern[ele - 'a'];
  }
  node->terminal = true;
}

vector<string> split(const string& sentence) {
  vector<string> word_vec;
  string pattern = "[[:alpha:]]*";
  regex r(pattern);
  for (sregex_iterator it(sentence.begin(), sentence.end(), r), end_it;
      it != end_it; ++it) {
    auto word = it->str();
    if (!word.empty())
      word_vec.push_back(it->str());
  }
  return word_vec;
}

string findPrefix(const string& word, PrefixNode* root) {
  string prefix;
  PrefixNode* node = root;
  for (auto ele : word) {
    if (node->terminal || node->childern[ele - 'a'] == nullptr) {
      break;
    }
    prefix.push_back(ele);
    node = node->childern[ele - 'a'];
  }
  return node->terminal ? prefix : "";
}

string replace(string sentence, vector<string> prefixs) {
  PrefixNode* root = new PrefixNode();
  for (auto ele : prefixs) {
    insertToPrefixTree(root, ele);
  }

  vector<string> sentence_vec = split(sentence);
  vector<string> replace_sentence_vec;
  for (auto ele : sentence_vec) {
    auto prefix = findPrefix(ele, root);
    if (prefix.empty()) {
      replace_sentence_vec.push_back(ele);
    } else {
      replace_sentence_vec.push_back(prefix);
    }
  }

  string replace_sentence;
  for (auto ele : replace_sentence_vec) {
    replace_sentence += (ele + " ");
  }
  return replace_sentence;
}

int main() {
  vector<string> prefixs = {"cat", "bat", "rat"};
  string sentence = "the cattle was rattled by the battery";
  string replace_sentence = replace(sentence, prefixs);
  cout << replace_sentence << endl;

  return 0;
}
